#include "../include/sm4.h"
#include <stdio.h>
#include <string.h>

void test_basic_encryption() {
    printf("\n=== SM4 基础加解密测试 ===\n");
    
    uint8_t key[SM4_KEY_SIZE];
    uint8_t plaintext[SM4_BLOCK_SIZE] = "20231310whb!!!!";
    uint8_t ciphertext[SM4_BLOCK_SIZE];
    uint8_t decrypted[SM4_BLOCK_SIZE];
    
    // 生成随机密钥
    sm4_generate_key(key);
    
    printf("测试数据: %s\n", plaintext);
    sm4_print_hex("随机密钥", key, SM4_KEY_SIZE);
    sm4_print_hex("明文", plaintext, SM4_BLOCK_SIZE);
    
    // 加密
    sm4_encrypt(key, plaintext, ciphertext);
    sm4_print_hex("密文", ciphertext, SM4_BLOCK_SIZE);
    
    // 解密
    sm4_decrypt(key, ciphertext, decrypted);
    sm4_print_hex("解密结果", decrypted, SM4_BLOCK_SIZE);
    
    // 验证
    if (memcmp(plaintext, decrypted, SM4_BLOCK_SIZE) == 0) {
        printf("✓ SM4 基础加解密测试成功！\n");
    } else {
        printf("✗ SM4 基础加解密测试失败！\n");
    }
}

void test_ecb_mode() {
    printf("\n=== SM4 ECB模式测试 ===\n");
    
    uint8_t key[SM4_KEY_SIZE];
    char message[] = "Hello SM4! This is ECB mode test for 20231310whb.";
    size_t msg_len = strlen(message);
    
    // 计算填充后的长度
    size_t padded_len = ((msg_len + SM4_BLOCK_SIZE) / SM4_BLOCK_SIZE) * SM4_BLOCK_SIZE;
    uint8_t ciphertext[256] = {0};
    uint8_t decrypted[256] = {0};
    
    // 生成随机密钥
    sm4_generate_key(key);
    
    printf("原始消息: %s\n", message);
    printf("消息长度: %zu\n", msg_len);
    printf("填充后长度: %zu\n", padded_len);
    sm4_print_hex("随机密钥", key, SM4_KEY_SIZE);
    sm4_print_hex("原始数据", (uint8_t*)message, msg_len);
    
    // ECB加密
    sm4_ecb_encrypt(key, (uint8_t*)message, ciphertext, msg_len);
    printf("ECB加密完成\n");
    sm4_print_hex("ECB密文", ciphertext, padded_len);
    
    // ECB解密
    sm4_ecb_decrypt(key, ciphertext, decrypted, padded_len);
    
    // 处理PKCS7填充
    size_t actual_len = padded_len;
    if (padded_len > 0) {
        uint8_t pad_value = decrypted[padded_len - 1];
        if (pad_value > 0 && pad_value <= SM4_BLOCK_SIZE) {
            // 验证填充是否正确
            int valid_padding = 1;
            for (size_t i = padded_len - pad_value; i < padded_len; i++) {
                if (decrypted[i] != pad_value) {
                    valid_padding = 0;
                    break;
                }
            }
            if (valid_padding) {
                actual_len = padded_len - pad_value;
            }
        }
    }
    
    // 确保字符串以null结尾
    decrypted[actual_len] = '\0';
    
    printf("解密后数据长度: %zu\n", actual_len);
    sm4_print_hex("解密数据", decrypted, actual_len);
    printf("ECB解密结果: %s\n", decrypted);
    
    if (actual_len == msg_len && memcmp(message, decrypted, msg_len) == 0) {
        printf("✓ SM4 ECB模式测试成功！\n");
    } else {
        printf("✗ SM4 ECB模式测试失败！\n");
        printf("  期望长度: %zu, 实际长度: %zu\n", msg_len, actual_len);
    }
}

int main() {
    printf("SM4 国密算法测试程序\n");
    printf("=====================\n");
    
    // 运行自检测试
    if (sm4_self_test() != 0) {
        return 1;
    }
    
    // 运行功能测试
    test_basic_encryption();
    test_ecb_mode();
    
    printf("\n=== 所有SM4测试完成 ===\n");
    return 0;
}
